{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([25, 10, 1])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "amp = torch.rand(25, 1) * 5\n",
    "phase = torch.rand(25, 1) * torch.pi\n",
    "init_inputs = (torch.rand(25, 10, 1) - 0.5) * 10\n",
    "amp.repeat_interleave(10, 1).unsqueeze(-1)\n",
    "outputs = amp.repeat_interleave(10, 1).unsqueeze(-1) * torch.sin(init_inputs - phase.repeat_interleave(10,1).unsqueeze(-1))\n",
    "outputs.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MAML(nn.Module):\n",
    "\n",
    "    def __init__(self,dim_input, dim_output, dim_hidden=[40, 40], *args, **kwargs) -> None:\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.dim_input = dim_input\n",
    "        self.dim_hidden = dim_hidden\n",
    "        self.dim_output = dim_output\n",
    "\n",
    "        self.weights = self.consturct_weights()\n",
    "\n",
    "    def consturct_weights(self):\n",
    "        weights = {}\n",
    "\n",
    "        weights['w1'] = nn.Parameter(torch.Tensor(self.dim_input, self.dim_hidden[0]).normal_(mean=0, std=0.01))\n",
    "        weights['b1'] = nn.Parameter(torch.zeros(self.dim_hidden[0]))\n",
    "        \n",
    "        for i in range(1, len(self.dim_hidden)):\n",
    "            weights['w'+str(i+1)] = nn.Parameter(torch.Tensor(self.dim_hidden[i-1], self.dim_hidden[i]).normal_(mean=0,std=0.01))\n",
    "            weights['b'+str(i+1)] = nn.Parameter(torch.zeros(self.dim_hidden[i]))\n",
    "        \n",
    "        weights['w'+str(len(self.dim_hidden)+1)] = nn.Parameter(torch.Tensor(self.dim_hidden[-1], self.dim_output).normal_(mean=0,std=0.01))\n",
    "        weights['b'+str(len(self.dim_hidden)+1)] = nn.Parameter(torch.zeros(self.dim_output))\n",
    "\n",
    "        return nn.ParameterDict(weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_points(k):\n",
    "    x = np.random.rand(k, 50)\n",
    "    y = np.random.choice([0,1], size=k, p=[.5, .5]).reshape([-1, 1])\n",
    "    return x, y\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.7396022  0.95965335 0.51709778 0.80024813 0.70184608 0.18649871\n",
      " 0.4741712  0.87541168 0.57172697 0.42414949 0.91501075 0.6388744\n",
      " 0.173118   0.34110602 0.30202243 0.8462335  0.70991907 0.93178554\n",
      " 0.08049162 0.4795771  0.27162653 0.96381474 0.67828474 0.32547353\n",
      " 0.90748402 0.09961256 0.76623766 0.34560302 0.64119806 0.51049803\n",
      " 0.93284522 0.49009912 0.17114519 0.70056469 0.08726381 0.08220708\n",
      " 0.23781147 0.39884318 0.43955274 0.91407789 0.3494363  0.04392989\n",
      " 0.02042178 0.10388845 0.17103966 0.3020508  0.65701021 0.76893113\n",
      " 0.03472028 0.76162472]\n",
      "[[0]\n",
      " [1]\n",
      " [0]\n",
      " [1]\n",
      " [0]\n",
      " [1]\n",
      " [1]\n",
      " [0]\n",
      " [1]\n",
      " [1]]\n"
     ]
    }
   ],
   "source": [
    "x, y = sample_points(10)\n",
    "print(x[0])\n",
    "print(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MAML:\n",
    "    def __init__(self):\n",
    "        self.num_tasks = 10\n",
    "        self.num_samples = 10\n",
    "        self.epochs = 10000  # \n",
    "        self.alpha = 0.0001  # inner loop\n",
    "        self.beta = 0.0001 # outer loop\n",
    "        self.theta = np.random.normal(size=50).reshape(50, 1)\n",
    "    \n",
    "    def sigmod(self, a):\n",
    "        return 1.0 / (1 + np.exp(-a))\n",
    "    \n",
    "    def train(self):\n",
    "        for e in range(self.epochs):\n",
    "            self.theta_ = []\n",
    "            for i in range(self.num_tasks):\n",
    "                X_train, Y_train = sample_points(self.num_samples)\n",
    "                a = np.matmul(X_train, self.theta)\n",
    "                YHat = self.sigmod(a)\n",
    "                loss = ((np.matmul(-Y_train.T, np.log(YHat)) - np.matmul((1 - Y_train.T), np.log(1 - YHat))) / self.num_samples)[0][0]\n",
    "                gradient = np.matmul(X_train.T, (YHat - Y_train)) / self.num_samples\n",
    "                self.theta_.append(self.theta - self.alpha * gradient)\n",
    "            meta_gradient = np.zeros(self.theta.shape)\n",
    "            for i in range(self.num_tasks):\n",
    "                XTest, YTest = sample_points(10)\n",
    "                a = np.matmul(XTest, self.theta_[i])\n",
    "                YPred = self.sigmod(a)\n",
    "                meta_gradient += np.matmul(XTest.T, (YPred - YTest)) / self.num_samples\n",
    "            self.theta = self.theta - self.beta * meta_gradient / self.num_tasks\n",
    "            if e % 1000 == 0:\n",
    "                print (\"Epoch {}: Loss {}\\n\".format(e,loss))\n",
    "                print ('Updated Model Parameter Theta\\n')\n",
    "                print ('Sampling Next Batch of Tasks \\n')\n",
    "                print ('---------------------------------\\n')\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0: Loss 2.186755963935155\n",
      "\n",
      "Updated Model Parameter Theta\n",
      "\n",
      "Sampling Next Batch of Tasks \n",
      "\n",
      "---------------------------------\n",
      "\n",
      "Epoch 1000: Loss 2.8275765667749004\n",
      "\n",
      "Updated Model Parameter Theta\n",
      "\n",
      "Sampling Next Batch of Tasks \n",
      "\n",
      "---------------------------------\n",
      "\n",
      "Epoch 2000: Loss 2.0792813463990876\n",
      "\n",
      "Updated Model Parameter Theta\n",
      "\n",
      "Sampling Next Batch of Tasks \n",
      "\n",
      "---------------------------------\n",
      "\n",
      "Epoch 3000: Loss 1.882167010856\n",
      "\n",
      "Updated Model Parameter Theta\n",
      "\n",
      "Sampling Next Batch of Tasks \n",
      "\n",
      "---------------------------------\n",
      "\n",
      "Epoch 4000: Loss 1.1809082982351349\n",
      "\n",
      "Updated Model Parameter Theta\n",
      "\n",
      "Sampling Next Batch of Tasks \n",
      "\n",
      "---------------------------------\n",
      "\n",
      "Epoch 5000: Loss 1.587422735155929\n",
      "\n",
      "Updated Model Parameter Theta\n",
      "\n",
      "Sampling Next Batch of Tasks \n",
      "\n",
      "---------------------------------\n",
      "\n",
      "Epoch 6000: Loss 2.047506563650271\n",
      "\n",
      "Updated Model Parameter Theta\n",
      "\n",
      "Sampling Next Batch of Tasks \n",
      "\n",
      "---------------------------------\n",
      "\n",
      "Epoch 7000: Loss 1.85590764536698\n",
      "\n",
      "Updated Model Parameter Theta\n",
      "\n",
      "Sampling Next Batch of Tasks \n",
      "\n",
      "---------------------------------\n",
      "\n",
      "Epoch 8000: Loss 0.9633381299747749\n",
      "\n",
      "Updated Model Parameter Theta\n",
      "\n",
      "Sampling Next Batch of Tasks \n",
      "\n",
      "---------------------------------\n",
      "\n",
      "Epoch 9000: Loss 1.0970365701706641\n",
      "\n",
      "Updated Model Parameter Theta\n",
      "\n",
      "Sampling Next Batch of Tasks \n",
      "\n",
      "---------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model = MAML()\n",
    "model.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ML",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
